from pyspark.sql import SparkSession
import os
import pyspark.sql.functions as F
from pyspark.sql.types import StringType

from cn.itcast.tag.base.BaseModel import BaseModel
from cn.itcast.tag.bean.ESMeta import ruleToESMeta

"""
-------------------------------------------------
   Description :	TODO：
   SourceFile  :	AgeModel1
   Author      :	itcast team
-------------------------------------------------
"""

# 0.设置系统环境变量
os.environ['JAVA_HOME'] = '/export/server/jdk1.8.0_241/'
os.environ['SPARK_HOME'] = '/export/server/spark'
os.environ['PYSPARK_PYTHON'] = '/root/anaconda3/envs/pyspark_env/bin/python3'
os.environ['PYSPARK_DRIVER_PYTHON'] = '/root/anaconda3/envs/pyspark_env/bin/python3'

#todo 年龄
class AgeModel1(BaseModel):
    def compute(self, es_df, five_df):
        """
        标签计算方法，根据业务数据和标签规则数据进行年龄段打标签操作

        Args:
            es_df: 从ES读取的业务数据，包含用户ID和生日信息
            five_df: 从MySQL读取的五级标签规则数据，包含年龄段划分规则

        Returns:
            DataFrame: 包含用户ID和对应标签ID的结果数据集
        """
        # 1. 清洗生日数据：去除生日中的"-"符号，以便转换为整数比较
        es_df = es_df.select("id", F.regexp_replace("birthday", "-", "").alias("birthday"))

        # 2. 解析年龄段规则：将rule字段按"-"拆分为起始年龄(start)和结束年龄(end)
        five_df = five_df.select(
            "id",
            F.split("rule", "-")[0].alias("start"),
            F.split("rule", "-")[1].alias("end")
        )

        # 3. 进行年龄段匹配：通过join操作将用户生日匹配到对应的年龄段规则
        new_df = es_df.join(
            other=five_df,
            on=es_df['birthday'].between(five_df['start'], five_df['end']),
            how='left'
        ).select(
            es_df['id'].cast(StringType()).alias("userId"),  # 用户ID
            five_df['id'].cast(StringType()).alias("tagsId")  # 匹配到的标签ID
        )

        return new_df


if __name__ == '__main__':
    ageModel = AgeModel1(14)
    ageModel.execute()
